import cmd
import sys
import os
import functools
import ollama
import json
import fitz  # PyMuPDF
import subprocess
import shlex
import asyncio
import traceback  # ### 改进点 1: 引入 traceback 模块 ###
from typing import List, Dict, Any, Optional, Callable


# 模拟 base_code 中的类，以便独立运行
class Tool:
    def __init__(self, name, description, function, required_args, optional_args=None):
        self.name = name
        self.description = description
        self.function = function
        self.required_args = required_args
        self.optional_args = optional_args or {}

    def __call__(self, **kwargs):
        return self.function(**kwargs)


class ToolSystem:
    def __init__(self):
        self._tools = {}

    def register_tool(self, tool: Tool):
        self._tools[tool.name] = tool

    def get_tool(self, name: str) -> Tool:
        return self._tools.get(name)

    def has_tool(self, name: str) -> bool:
        return name in self._tools

    def list_tools(self) -> List[str]:
        return list(self._tools.keys())


# ==============================================================================
# 1. CommandRegistry
# ==============================================================================
class CommandRegistry(ToolSystem):
    def __init__(self):
        super().__init__()
        self.lines = [
            "     欢迎使用 SimpleMAS ReAct v1.2 (ReAct 模式)",
            " ",
            "支持自然语言交互，Agent 将通过 思考->行动->观察 循环解决复杂问题",
            " ",
            "您可以直接下达复杂指令, 例如: ",
            "  '请列出当前目录下的所有pdf文件，然后读取 paper.pdf 的第一页并总结其内容'",
            "  '进入 src 目录，然后列出所有文件'",
            " ",
            "内置命令:",
            "  list_tools             - 列出所有已注册的工具及其详细描述",
            "  help                   - 显示此帮助信息",
            "  exit / quit            - 退出程序"
        ]

    def generate_tools_description(self) -> str:
        if not self._tools:
            return "没有可用的工具。"
        descriptions = []
        for tool_name, tool in self._tools.items():
            params = []
            for arg in tool.required_args:
                if arg != 'memory':
                    params.append(f"{arg}: str")
            if tool.optional_args is not None:
                for arg, default in tool.optional_args.items():
                    if arg != 'memory':
                        params.append(f"{arg}: {type(default).__name__} = {repr(default)}")
            param_str = ", ".join(params)
            descriptions.append(f"- {tool.name}({param_str}): {tool.description}")
        return "\n".join(descriptions)


GLOBAL_TOOL_SYSTEM = CommandRegistry()


def tool(name: str, description: str, required_args: List[str], optional_args: Optional[Dict[str, Any]] = None):
    def decorator(func: Callable):
        new_tool = Tool(
            name=name,
            description=description,
            function=func,
            required_args=required_args,
            optional_args=optional_args
        )
        GLOBAL_TOOL_SYSTEM.register_tool(new_tool)

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)

        return wrapper

    return decorator


# ==============================================================================
# 2. HelpAgent
# ==============================================================================
class HelpAgent:
    def __init__(self, registry: CommandRegistry):
        self.command_registry = registry

    def show_intro(self):
        return "\n".join(self.command_registry.lines)

    def show_help(self, command=None):
        if command and self.command_registry.has_tool(command):
            t = self.command_registry.get_tool(command)
            usage = f"{t.name} {' '.join(arg for arg in t.required_args if arg != 'memory')}"
            return f"命令: {t.name}\n功能: {t.description}\n用法: {usage}"
        else:
            return self.show_intro()


# ==============================================================================
# 3. AgentMemory
# ==============================================================================
class AgentMemory:
    def __init__(self):
        self.beliefs = {}


# ==============================================================================
#  ReActAgent
# ==============================================================================
class ReActAgent:
    def __init__(self, tool_system: CommandRegistry,
                 model_name: str = "qwen3:4b",
                 max_iterations: int = 10):
        self.tool_system = tool_system
        self.model_name = model_name
        self.max_iterations = max_iterations
        self.ollama_client = ollama.AsyncClient()
        self.memory = AgentMemory()
        # ### 改进点 2: 增加调试开关 ###
        # 如果需要查看发送给LLM的完整提示，请将此项设为 True
        self.debug_mode = False

    def _create_react_prompt(self, user_query: str, history: List[str]) -> str:
        tools_description = self.tool_system.generate_tools_description()
        scratchpad = "\n".join(history) if history else "无"
        return f"""
        你是一个能够解决复杂问题的智能助手。你的任务是接收用户的请求，并通过一个“思考->行动->观察”的循环来完成它。
        你必须严格遵守以下规则和工作流程。

        **可用工具:**
        你只能使用以下定义的工具。不要编造任何不存在的工具。
        {tools_description}

        **重要规则:**
        1.  **语言一致性**: 你的"思考"和"最终答案"必须使用与用户请求相同的语言。如果用户用中文提问，你就必须用中文回答。
        2.  **当前目录**: 所有文件操作和终端命令都相对于当前工作目录执行。当前目录是 `{os.getcwd()}`。

        **工作流程:**
        1.  **思考 (Thought)**: 分析用户请求和历史记录，制定下一步计划。决定是调用工具 (`tool_call`) 还是给出最终答案 (`final_answer`)。
        2.  **行动 (Action)**: 如果需要调用工具，你的输出必须是一个严格的JSON对象，格式如下：
            ```json
            {{
                "thought": "这里是你的中文思考过程...",
                "tool_call": {{
                    "name": "工具名称",
                    "arguments": {{
                        "参数1": "值1"
                    }}
                }}
            }}
            ```
        3.  **最终答案 (Final Answer)**: 如果任务已完成，你的输出必须是另一个严格的JSON对象，格式如下：
            ```json
            {{
                "thought": "我已经完成了所有步骤，现在可以给出最终答案了。",
                "final_answer": "这里是给用户的最终回答（必须使用用户提问的语言）。"
            }}
            ```

        **当前任务**
        用户请求: "{user_query}"
        历史记录 (已完成的步骤):
        {scratchpad}

        请严格按照JSON格式输出你的下一步。
        """

    async def run_async(self, user_query: str) -> str:
        history = []
        for i in range(self.max_iterations):
            # ### 改进点 3: 强制刷新输出 ###
            print(f"\n🌀 === [循环 {i + 1}/{self.max_iterations}] ===", flush=True)

            system_prompt = self._create_react_prompt(user_query, history)

            if self.debug_mode:
                print("\n--- DEBUG: PROMPT SENT TO LLM ---\n", flush=True)
                print(system_prompt, flush=True)
                print("\n--- END OF DEBUG PROMPT ---\n", flush=True)

            print("🧠 正在思考下一步...", flush=True)
            content = ""
            try:
                response = await self.ollama_client.chat(
                    model=self.model_name,
                    messages=[{"role": "system", "content": system_prompt}],
                    format="json"
                )
                content = response["message"]["content"]
                llm_output = json.loads(content)

                thought = llm_output.get("thought", "(无思考)")
                print(f"🤔 思考: {thought}", flush=True)

            except Exception as e:
                # ### 改进点 1: 强化错误诊断 ###
                print("\n" + "=" * 20 + " ❌ ReAct 循环内部错误 " + "=" * 20, flush=True)
                print(f"在第 {i + 1} 次循环中，LLM响应或JSON解析失败。", flush=True)
                print(f"错误类型: {type(e).__name__}", flush=True)
                print(f"错误详情: {e}", flush=True)
                print(f"LLM原始响应内容 (可能不完整): \n---\n{content}\n---", flush=True)
                print("详细错误追溯:", flush=True)
                traceback.print_exc()  # 打印完整的错误栈
                print("=" * 60, flush=True)

                error_msg = f"LLM响应错误或JSON解析失败: {e}"
                history.append(f"Observation: {error_msg}")
                continue

            if "final_answer" in llm_output:
                final_answer = llm_output["final_answer"]
                print(f"✅ 任务完成。", flush=True)
                return final_answer

            if "tool_call" in llm_output:
                tool_call = llm_output["tool_call"]
                tool_name = tool_call.get("name")
                tool_args = tool_call.get("arguments", {})

                if not tool_name:
                    observation = "错误: LLM的工具调用缺少 'name' 字段。"
                elif not self.tool_system.has_tool(tool_name):
                    observation = f"错误: 尝试调用一个不存在的工具 '{tool_name}'。可用工具: {self.tool_system.list_tools()}"
                else:
                    print(f"🛠️ 行动: 调用工具 `{tool_name}`，参数: {tool_args}", flush=True)
                    try:
                        tool_obj = self.tool_system.get_tool(tool_name)
                        if 'memory' in tool_obj.required_args:
                            tool_args["memory"] = self.memory

                        maybe_coro = tool_obj(**tool_args)
                        if asyncio.iscoroutine(maybe_coro):
                            observation = await maybe_coro
                        else:
                            observation = maybe_coro

                        observation = str(observation)

                        if len(observation) > 1500:
                            observation = observation[:1500] + "\n... (内容已截断)"

                    except Exception as e:
                        observation = f"❌ 工具 '{tool_name}' 执行失败: {type(e).__name__}: {e}"

                print(f"👀 观察: {observation}", flush=True)
                history.append(f"Thought: {thought}")
                history.append(f"Action: 调用 {tool_name}({tool_args})")
                history.append(f"Observation: {observation}")
            else:
                observation = "错误: LLM的响应格式不正确，既没有 final_answer 也没有 tool_call。"
                print(f"👀 观察: {observation}", flush=True)
                history.append(f"Observation: {observation}")

        return "🔁 Agent 已达到最大循环次数，但仍未得出最终答案。请尝试更具体地提问或分解任务。"


# ==============================================================================
# 注册工具
# ==============================================================================
@tool("change_directory", "更改当前工作目录", ["path"])
def change_directory(path: str) -> str:
    try:
        os.chdir(path)
        new_cwd = os.getcwd()
        return f"✅ 成功将目录更改为: {new_cwd}"
    except FileNotFoundError:
        return f"❌ 错误: 目录 '{path}' 不存在。"
    except Exception as e:
        return f"❌ 错误: 更改目录时发生未知错误: {e}"


import locale  # 导入 locale 模块


@tool("run_terminal_command", "执行一个非交互式的系统终端命令 (例如 ls, dir, cat, echo)。禁止使用此工具执行 'cd'。",
      ["command_line"])
def run_terminal_command(command_line: str) -> str:
    """执行一个完整的命令行字符串，并智能处理不同操作系统的输出编码。"""
    command_line = command_line.strip()
    if not command_line:
        return "❌ 错误: 未提供任何命令。"

    if command_line.startswith("cd "):
        return "❌ 错误: 禁止使用 'run_terminal_command' 来执行 'cd'。请使用 'change_directory' 工具来更改目录。"

    try:
        # 1. 移除 text=True 和 encoding，捕获原始字节流
        result = subprocess.run(
            command_line,
            shell=True,
            capture_output=True,
            cwd=os.getcwd()
        )

        # 2. 获取系统首选编码 (例如在中文Windows上是 'gbk')
        #    这是解决 UnicodeDecodeError 的关键
        encoding = locale.getpreferredencoding(False)

        # 3. 使用正确的编码解码 stdout 和 stderr 的字节流
        #    使用 errors='replace' 确保即使有无法解码的字符程序也不会崩溃
        output = result.stdout.decode(encoding, errors='replace').strip() if result.stdout else ""
        error = result.stderr.decode(encoding, errors='replace').strip() if result.stderr else ""

        if result.returncode == 0:
            return f"✅ 命令成功执行。\n输出:\n---\n{output or '(无输出)'}\n---"
        else:
            # 同时返回标准输出和标准错误，给Agent更完整的上下文
            return f"❌ 命令执行失败 (返回码: {result.returncode})。\n错误信息:\n---\n{error or '(无错误信息)'}\n---\n输出:\n---\n{output or '(无输出)'}\n---"

    except Exception as e:
        # 捕获其他意外异常，例如命令本身不存在
        return f"❌ 执行系统命令时发生异常: {type(e).__name__}: {e}"


@tool("read_pdf", "读取PDF文件前N页的文本内容", ["path"], {"max_pages": 1})
def read_pdf_text(path: str, max_pages: int = 1) -> str:
    if not os.path.exists(path):
        return f"❌ 错误: 文件 '{path}' 未找到。当前目录是 '{os.getcwd()}'。"
    if not path.lower().endswith('.pdf'):
        return "❌ 错误: 文件必须是PDF格式。"
    text = []
    try:
        with fitz.open(path) as doc:
            for page_num in range(min(doc.page_count, max_pages)):
                page = doc.load_page(page_num)
                text.append(page.get_text("text"))
        return "\n".join(text) if text else "PDF页面中未找到可读文本。"
    except Exception as e:
        return f"❌ 读取PDF时发生错误: {e}"


@tool("search_pdf", "在PDF文件中搜索关键字", ["path", "keyword"])
def search_pdf(path: str, keyword: str) -> str:
    if not os.path.exists(path):
        return f"❌ 错误: 文件 '{path}' 未找到。当前目录是 '{os.getcwd()}'。"
    results = []
    with fitz.open(path) as doc:
        for page_num, page in enumerate(doc, start=1):
            instances = page.search_for(keyword, quads=True)
            if instances:
                for inst in instances:
                    context_rect = inst.irect.inflate(20, 20)
                    context = page.get_text("text", clip=context_rect).strip().replace("\n", " ")
                    results.append({"page": page_num, "context": f"...{context}..."})
    if not results:
        return f"在文件 '{path}' 中未找到关键字 '{keyword}'。"
    return json.dumps(results, ensure_ascii=False, indent=2)


# =======================================================================
# SimpleAgentShell
# =======================================================================
class SimpleAgentShell(cmd.Cmd):
    help_agent = HelpAgent(GLOBAL_TOOL_SYSTEM)
    cmd_agent = ReActAgent(GLOBAL_TOOL_SYSTEM)
    intro = help_agent.show_intro()

    @property
    def prompt(self):
        return f"[{os.path.basename(os.getcwd())}] agent> "

    def __init__(self):
        super().__init__()

    def onecmd(self, line):
        """
        cmd.Cmd 类的 onecmd 方法是所有命令分发的中心枢纽
        - 在 do_<command> 被调用之前，onecmd 会接收到完整的用户输入行
        - 如果输入以 / 开头: 我们会将其视为一个直接命令。我们会去掉 /，然后
          调用 super().onecmd()，让 cmd 模块按其标准流程处理这个命令（例如，
          '/help' 会被分派到 do_help 方法）。
        - 如果输入不以 / 开头: 我们会将其视为一个自然语言任务，并将其传递给
          ReActAgent 进行处理。这部分逻辑将从原来的 default 方法移到 onecmd 中。
        """
        line = line.strip()
        if not line:
            return self.emptyline()  # 处理空行

        if line.startswith('/'):
            # 这是直接命令，去掉'/'后交由cmd框架的标准流程处理
            return super().onecmd(line[1:])
        else:
            # 这是自然语言任务，交给Agent处理
            print("🚀 检测到自然语言任务，正在启动 Agent...")
            try:
                final_result = asyncio.run(self.cmd_agent.run_async(line))
                print("\n" + "=" * 20 + " Agent 最终答复 " + "=" * 20, flush=True)
                print(final_result, flush=True)
                print("=" * 55, flush=True)
            except Exception as e:
                print(f"\n❌ Agent在执行过程中遇到严重错误: {e}", flush=True)
                traceback.print_exc()
            return False  # 返回False以继续等待下一条命令

    def default(self, line):
        """
        - 未知的直接命令时被调用（例如，用户输入了 /non_existent_command）。
        - 我们可以让它打印一个更清晰的错误提示。
        """
        print(f"❌ 未知直接命令: '/{line}'.")
        print("   请输入 '/help' 查看所有可用的直接命令。")
        print("   对于自然语言任务，请直接输入，不要加'/'前缀。")

    def do_list_tools(self, arg):
        print("\n--- 所有可用工具 ---")
        print(self.cmd_agent.tool_system.generate_tools_description())
        print("--------------------\n")

    def do_exit(self, arg):
        print("\n👋 感谢使用，再见！")
        return True

    def do_quit(self, arg):
        return self.do_exit(arg)

    def do_EOF(self, arg):
        print(); return self.do_exit(arg)

    def do_help(self, arg):
        print(self.help_agent.show_help(arg.strip()))


def main():
    try:
        SimpleAgentShell().cmdloop()
    except KeyboardInterrupt:
        print("\n👋 检测到中断，程序退出。")
        sys.exit(0)


if __name__ == '__main__':
    main()